package com.cloud.common.support;

import com.cloud.common.bean.SessionInfo;
import com.cloud.common.context.SessionContext;
import com.cloud.common.context.TraceContext;
import com.cloud.common.context.VersionContext;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.*;

/**
 * 默认线程池
 */
@Slf4j
public class DefaultThreadPoolExecutor extends ThreadPoolExecutor {

    public DefaultThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public DefaultThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    public DefaultThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
    }

    public DefaultThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    @Override
    public void execute(Runnable task) {
        String traceId = TraceContext.get();
        String version = VersionContext.get();
        SessionInfo sessionInfo = SessionContext.get();
        super.execute(() -> run(task, traceId, version, sessionInfo));
    }

    @Override
    public Future<?> submit(Runnable task) {
        String traceId = TraceContext.get();
        String version = VersionContext.get();
        SessionInfo sessionInfo = SessionContext.get();
        return super.submit(() -> run(task, traceId, version, sessionInfo));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        String traceId = TraceContext.get();
        String version = VersionContext.get();
        SessionInfo sessionInfo = SessionContext.get();
        return super.submit(() -> run(task, traceId, version, sessionInfo));
    }

    @Override
    public <T> Future<T> submit(Runnable task, T result) {
        String traceId = TraceContext.get();
        String version = VersionContext.get();
        SessionInfo sessionInfo = SessionContext.get();
        return super.submit(() -> run(task, traceId, version, sessionInfo), result);
    }

    private void run(Runnable task, String traceId, String version, SessionInfo sessionInfo) {
        try {
            TraceContext.set(traceId);
            VersionContext.set(version);
            SessionContext.set(sessionInfo);
            task.run();
        } finally {
            TraceContext.remove();
        }
    }

    private <T> T run(Callable<T> task, String traceId, String version, SessionInfo sessionInfo) throws Exception {
        try {
            TraceContext.set(traceId);
            VersionContext.set(version);
            SessionContext.set(sessionInfo);
            return task.call();
        } finally {
            TraceContext.remove();
        }
    }
}
